package org.example.leetCode;

/**
 * 215. 数组中的第K个最大元素
 * 给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。
 * <p>
 * 请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。
 * <p>
 * 你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。
 */
public class Code215 {
    /**
     * 使用快速选择算法找到第k个最大元素
     * 以最左边的元素为基准元素
     *
     * @param nums 整数数组
     * @param k 第k个最大元素
     * @return 第k个最大元素的值
     */
    public int findKthLargest(int[] nums, int k) {
        return quickSelect(nums, 0, nums.length - 1, k);
    }

    /**
     * 快速选择算法主函数
     * @param nums 数组
     * @param left 左边界
     * @param right 右边界
     * @param k 第k个最大元素
     * @return 第k个最大元素
     */
    private int quickSelect(int[] nums, int left, int right, int k) {
        if (left == right) {
            return nums[left];
        }

        // 进行分区操作，返回基准元素的最终位置
        int pivotIndex = partition(nums, left, right);

        // pivotIndex是基准元素在排序后数组中的位置（从0开始）
        // 转换为"第几大"需要: nums.length - pivotIndex
        int rank = pivotIndex - left + 1; // 基准元素是当前子数组中第rank大的元素

        if (rank == k) {
            // 找到第k大的元素
            return nums[pivotIndex];
        } else if (rank > k) {
            // 第k大的元素在左半部分
            return quickSelect(nums, left, pivotIndex - 1, k);
        } else {
            // 第k大的元素在右半部分
            return quickSelect(nums, pivotIndex + 1, right, k - rank);
        }
    }

    /**
     * 分区操作：以最左边元素为基准，将数组分为大于基准和小于基准的两部分
     * @param nums 数组
     * @param left 左边界
     * @param right 右边界
     * @return 基准元素的最终位置
     */
    private int partition(int[] nums, int left, int right) {
        // 以最左边的元素为基准
        int pivot = nums[left];
        int i = left + 1;  // 从基准元素的下一个位置开始
        int j = right;     // 从最右边开始

        while (true) {
            // 从左边找到第一个小于等于基准的元素
            while (i <= j && nums[i] > pivot) {
                i++;
            }

            // 从右边找到第一个大于基准的元素
            while (i <= j && nums[j] <= pivot) {
                j--;
            }

            // 如果指针相遇，则停止
            if (i >= j) {
                break;
            }

            // 交换元素
            swap(nums, i, j);
            i++;
            j--;
        }

        // 将基准元素放到正确位置
        swap(nums, left, j);
        return j;
    }

    /**
     * 基于你原始思路的实现（修正版）
     * 以最左边元素为基准元素
     */
    public int findKthLargestOriginal(int[] nums, int k) {
        return quickSelectOriginal(nums, 0, nums.length - 1, k);
    }

    private int quickSelectOriginal(int[] nums, int left, int right, int k) {
        if (left >= right) {
            return nums[left];
        }

        int pivot = nums[left];  // 基准元素
        int i = left + 1;        // 左指针
        int j = right;           // 右指针

        while (i <= j) {
            // 从左边找到第一个小于等于基准的元素
            while (i <= right && nums[i] > pivot) {
                i++;
            }

            // 从右边找到第一个大于基准的元素
            while (j > left && nums[j] <= pivot) {
                j--;
            }

            // 如果i < j，交换元素
            if (i < j) {
                swap(nums, i, j);
                i++;
                j--;
            } else {
                break;
            }
        }

        // 将基准元素放到正确位置
        swap(nums, left, j);

        // 计算基准元素是当前范围中第几大的元素
        int rank = j - left + 1;

        if (rank == k) {
            return nums[j];
        } else if (rank > k) {
            // 第k大的元素在左半部分
            return quickSelectOriginal(nums, left, j - 1, k);
        } else {
            // 第k大的元素在右半部分
            return quickSelectOriginal(nums, j + 1, right, k - rank);
        }
    }

    /**
     * 交换数组中两个元素的位置
     * @param nums 数组
     * @param i 索引i
     * @param j 索引j
     */
    private void swap(int[] nums, int i, int j) {
        int temp = nums[i];
        nums[i] = nums[j];
        nums[j] = temp;
    }

    public static void main(String[] args) {
        Code215 solution = new Code215();

        // 测试用例1
        int[] nums1 = {3, 2, 1, 5, 6, 4};
        int k1 = 2;
        int result1 = solution.findKthLargest(nums1, k1);
        System.out.println("数组: [3,2,1,5,6,4], k=2");
        System.out.println("第" + k1 + "个最大元素: " + result1);
        System.out.println("预期结果: 5");
        System.out.println();

        // 测试用例2
        int[] nums2 = {3, 2, 3, 1, 2, 4, 5, 5, 6};
        int k2 = 4;
        int result2 = solution.findKthLargest(nums2, k2);
        System.out.println("数组: [3,2,3,1,2,4,5,5,6], k=4");
        System.out.println("第" + k2 + "个最大元素: " + result2);
        System.out.println("预期结果: 4");
        System.out.println();

        // 测试用例3
        int[] nums3 = {1};
        int k3 = 1;
        int result3 = solution.findKthLargest(nums3, k3);
        System.out.println("数组: [1], k=1");
        System.out.println("第" + k3 + "个最大元素: " + result3);
        System.out.println("预期结果: 1");
        System.out.println();

        // 测试原始方法
        int result4 = solution.findKthLargestOriginal(nums1, k1);
        System.out.println("使用原始方法测试用例1结果: " + result4);
    }
}
